#!/usr/bin/env python
# -*- coding: UTF-8 -*-
"""
@Project ：python_learning 
@File ：MySQLORMGenerator.py
@IDE  ：PyCharm 
@Author ：李涵彬
@Date ：2025/7/9 下午3:32 
"""

# !/usr/bin/env python
# -*- coding: UTF-8 -*-
"""
@Project ：python_learning
@File ：MySQLORMGenerator.py
@IDE  ：PyCharm
@Author ：李涵彬
@Date ：2025/7/9
"""

import pymysql
import os
import logging
from datetime import datetime
from typing import List, Dict

# 配置日志
logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %(message)s')
logger = logging.getLogger(__name__)

# 数据库配置
MYSQL_DATABASE = 'sgjt_new'
MYSQL_HOST = '127.0.0.1'
MYSQL_PORT = 3306
MYSQL_USER = 'root'
MYSQL_PASSWORD = '12345'
MYSQL_CHARSET = 'utf8mb4'

# 字段名映射字典（中文 -> 英文）
FIELD_NAME_MAPPING = {
	"序号": "id",
	"产线": "production_line",
	"字段": "field_name",
	"操作": "operation",
	"类型": "type",
	"值一": "value1",
	"值二": "value2",
	# 你可以继续添加更多映射
}


class MySQLORMGenerator:
	def __init__(self):
		self.connection = None
		self.output_dir = "mysql_models"
		self.type_mapping = {
			'int': 'int',
			'bigint': 'int',
			'smallint': 'int',
			'tinyint': 'int',
			'bit': 'bool',
			'decimal': 'float',
			'numeric': 'float',
			'float': 'float',
			'double': 'float',
			'date': 'datetime.date',
			'time': 'datetime.time',
			'datetime': 'datetime.datetime',
			'timestamp': 'datetime.datetime',
			'char': 'str',
			'varchar': 'str',
			'text': 'str',
			'longtext': 'str',
			'mediumtext': 'str',
			'tinytext': 'str',
			'json': 'dict',
			'blob': 'bytes',
			'mediumblob': 'bytes',
			'longblob': 'bytes',
			'binary': 'bytes',
			'varbinary': 'bytes',
			'enum': 'str',
			'set': 'str'
		}

	def connect(self) -> bool:
		try:
			self.connection = pymysql.connect(
				host=MYSQL_HOST,
				port=MYSQL_PORT,
				user=MYSQL_USER,
				password=MYSQL_PASSWORD,
				database=MYSQL_DATABASE,
				charset=MYSQL_CHARSET
			)
			logger.info("MySQL 数据库连接成功")
			return True
		except Exception as e:
			logger.error(f"数据库连接失败: {e}")
			return False

	def get_all_tables(self) -> List[str]:
		try:
			with self.connection.cursor() as cursor:
				cursor.execute("SHOW TABLES")
				tables = [row[0] for row in cursor.fetchall()]
				logger.info(f"找到 {len(tables)} 个表")
				return tables
		except Exception as e:
			logger.error(f"获取表列表失败: {e}")
			return []

	def get_table_columns(self, table_name: str) -> List[Dict]:
		try:
			with self.connection.cursor() as cursor:
				cursor.execute(f"DESCRIBE `{table_name}`")
				columns = []
				for row in cursor.fetchall():
					field, type_str, null, key, default, extra = row
					columns.append({
						'original_name': field,
						'name': FIELD_NAME_MAPPING.get(field, field),  # 使用映射名
						'type': type_str.split('(')[0],
						'nullable': null == 'YES',
						'default': default,
						'key': key,
						'extra': extra
					})
				return columns
		except Exception as e:
			logger.error(f"获取表 {table_name} 的字段信息失败: {e}")
			return []

	def get_foreign_keys(self, table_name: str) -> List[Dict]:
		try:
			with self.connection.cursor() as cursor:
				cursor.execute("""
                    SELECT
                        COLUMN_NAME,
                        REFERENCED_TABLE_NAME,
                        REFERENCED_COLUMN_NAME,
                        CONSTRAINT_NAME
                    FROM INFORMATION_SCHEMA.KEY_COLUMN_USAGE
                    WHERE TABLE_SCHEMA = %s AND TABLE_NAME = %s AND REFERENCED_TABLE_NAME IS NOT NULL
                """, (MYSQL_DATABASE, table_name))
				foreign_keys = []
				for row in cursor.fetchall():
					foreign_keys.append({
						'column': FIELD_NAME_MAPPING.get(row[0], row[0]),
						'referenced_table': row[1],
						'referenced_column': row[2],
						'constraint_name': row[3]
					})
				return foreign_keys
		except Exception as e:
			logger.error(f"获取表 {table_name} 的外键失败: {e}")
			return []

	def convert_type(self, sql_type: str, nullable: bool) -> str:
		python_type = self.type_mapping.get(sql_type.lower(), 'str')
		if nullable and python_type not in ['object', 'dict']:
			return f"Optional[{python_type}]"
		return python_type

	def snake_to_pascal(self, snake_str: str) -> str:
		return ''.join(word.capitalize() for word in snake_str.split('_'))

	def generate_model_class(self, table_name: str) -> str:
		columns = self.get_table_columns(table_name)
		foreign_keys = self.get_foreign_keys(table_name)
		primary_keys = [col['name'] for col in columns if col['key'] == 'PRI']

		if not columns:
			return ""

		class_name = self.snake_to_pascal(table_name)
		code_lines = []

		code_lines.append(f'"""')
		code_lines.append(f'自动生成的ORM模型类 - {table_name}')
		code_lines.append(f'生成时间: {datetime.now().strftime("%Y-%m-%d %H:%M:%S")}')
		code_lines.append('"""')
		code_lines.append('')

		code_lines.append('from dataclasses import dataclass')
		code_lines.append('from typing import Optional')
		code_lines.append('import datetime')
		code_lines.append('')

		code_lines.append('@dataclass')
		code_lines.append(f'class {class_name}:')
		code_lines.append(f'    """{table_name} 表的 ORM 模型类"""')
		code_lines.append('')
		code_lines.append(f'    __tablename__ = "{table_name}"')
		code_lines.append('')

		if primary_keys:
			code_lines.append(f'    # 主键: {", ".join(primary_keys)}')

		if foreign_keys:
			code_lines.append('    # 外键关系:')
			for fk in foreign_keys:
				code_lines.append(f'    # {fk["column"]} -> {fk["referenced_table"]}.{fk["referenced_column"]}')

		code_lines.append('')

		for col in columns:
			comment_parts = [f"原字段: {col['original_name']} | SQL类型: {col['type']}"]
			if not col['nullable']:
				comment_parts.append("NOT NULL")
			if col['default']:
				comment_parts.append(f"默认值: {col['default']}")
			if col['name'] in primary_keys:
				comment_parts.append("主键")
			for fk in foreign_keys:
				if fk['column'] == col['name']:
					comment_parts.append(f"外键 -> {fk['referenced_table']}.{fk['referenced_column']}")

			comment = " | ".join(comment_parts)
			python_type = self.convert_type(col['type'], col['nullable'])
			code_lines.append(f'    {col["name"]}: {python_type}  # {comment}')

		code_lines.append('')

		code_lines.append('    def to_dict(self) -> dict:')
		code_lines.append('        return {')
		for col in columns:
			code_lines.append(f'            "{col["name"]}": self.{col["name"]},')
		code_lines.append('        }')

		return '\n'.join(code_lines)

	def generate_init_file(self, table_names: List[str]) -> str:
		code_lines = ['"""', '自动生成的模型包初始化文件', f'生成时间: {datetime.now().strftime("%Y-%m-%d %H:%M:%S")}',
					  '"""', '']
		for table in sorted(table_names):
			class_name = self.snake_to_pascal(table)
			code_lines.append(f'from .{table} import {class_name}')
		code_lines.append('')
		code_lines.append('__all__ = [')
		for table in sorted(table_names):
			class_name = self.snake_to_pascal(table)
			code_lines.append(f'    "{class_name}",')
		code_lines.append(']')
		return '\n'.join(code_lines)

	def generate_models(self):
		if not self.connect():
			return

		if not os.path.exists(self.output_dir):
			os.makedirs(self.output_dir)

		tables = self.get_all_tables()
		successful = []

		for table in tables:
			try:
				logger.info(f"正在生成 {table} 的 ORM 类...")
				model_code = self.generate_model_class(table)
				if model_code:
					with open(os.path.join(self.output_dir, f"{table}.py"), 'w', encoding='utf-8') as f:
						f.write(model_code)
					successful.append(table)
					logger.info(f"✓ {table} 生成成功")
				else:
					logger.warning(f"✗ {table} 生成失败")
			except Exception as e:
				logger.error(f"生成 {table} 时出错: {e}")

		init_code = self.generate_init_file(successful)
		with open(os.path.join(self.output_dir, "__init__.py"), 'w', encoding='utf-8') as f:
			f.write(init_code)
		logger.info("✓ __init__.py 生成成功")

		logger.info("=" * 60)
		logger.info(f"成功生成 {len(successful)} 个模型类")


def main():
	print("MySQL ORM 模型生成器")
	print("=" * 50)
	generator = MySQLORMGenerator()
	generator.generate_models()


if __name__ == "__main__":
	main()
